﻿using FaissOperation.FaissIndex.Index;
using FaissOperation.Interface;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace FaissOperation.FaissIndex
{
    public abstract class BaseFaissCollection : IVectorCollection<float[]>
    {
        internal const int DEFAULT_DIMENSION = 1024; // 默认维度1024

        private int _dimension = DEFAULT_DIMENSION;
        internal IFaissIndex faissIndex;

        internal Dictionary<int, float[]> vectors = new Dictionary<int, float[]>();

        public int Dimension
        {
            get { return _dimension; }
            set
            {
                if (vectors != null && vectors.Count > 0)
                    throw new InvalidOperationException("在集合中有向量存在的情况下，不允许修改维度。");

                _dimension = value;
                CreateIndex();
            }
        }

        public bool NormalizeOnAddAndQuery { get; set; } = false;

        public MetricType MetricTypeInUse { get; set; } = MetricType.L2;

        public void AddVector(int id, float[] vector)
        {
            var v = NormalizeOnAddAndQuery ? FaissHelper.NormalizeForFaiss(vector) : vector;
            vectors[id] = v;
            faissIndex.Add(v);
        }

        public void AddVectors(IEnumerable<(int id, float[] Vector)> items)
        {
            var vectors = new List<float[]>();
            foreach (var (id, vector) in items)
            {
                var v = NormalizeOnAddAndQuery ? FaissHelper.NormalizeForFaiss(vector) : vector;
                this.vectors[id] = v;
                vectors.Add(v);
            }

            faissIndex.BatchAdd(vectors.ToArray());
        }

        public void RemoveVector(int id)
        {
            if (vectors.Remove(id))
                RebuildIndex();
        }

        public void RemoveVectors(IEnumerable<int> ids)
        {
            bool removed = false;
            foreach (var id in ids)
            {
                if (vectors.Remove(id)) removed = true;
            }
            if (removed)
                RebuildIndex();
        }
        public void Clear()
        {
            vectors.Clear();
            RebuildIndex();
        }

        public void LoadData(PlainCollectionData data)
        {
            FaissHelper.AssertFaissDeployment();

            vectors.Clear();

            this.Dimension = data.Dimension;
            this.MetricTypeInUse = data.MetricType;
            this.NormalizeOnAddAndQuery = data.NormalizeOnAdd;

            for (int i = 0; i < data.Vectors.Count; i++)
            {
                vectors.Add(data.Ids[i], data.Vectors[i]);
            }

            RebuildIndex();
        }

        public PlainCollectionData GetData()
        {
            return new PlainCollectionData
            {
                Dimension = this.Dimension,
                MetricType = this.MetricTypeInUse,
                NormalizeOnAdd = this.NormalizeOnAddAndQuery,
                Ids = vectors.Keys.ToList(),
                Vectors = vectors.Values.ToList()
            };
        }

        public long Count => vectors.Count;

        public void Destory()
        {
            RebuildIndex();
            this.faissIndex.Dispose();
        }

        public abstract IEnumerable<SearchResult> Search(float[] query, int topK);

        /// <summary>
        /// 重建索引，通常在添加或删除向量后调用。
        /// </summary>
        public void RebuildIndex()
        {
            faissIndex.Rebuild();
            if (vectors.Count > 0)
            {
                faissIndex.BatchAdd(vectors.Values.ToArray());
            }
        }


        /// <summary>
        /// 创建索引，通常在设置维度后调用。
        /// 根据具体的索引类型实现不同的索引创建逻辑。如 IVF、Flat 等。
        /// </summary>
        public abstract void CreateIndex();


        public void Train(float[][] vectors)
        {
            faissIndex.Train(vectors);
        }

        public bool IsTrained => faissIndex.IsTrained;



    }
}
